/**
 * @file test_pdmat.cpp
 *
 * Unit testing of pdmat classes
 *
 * @author Dahua Lin
 */

#include "test_common.h"
#include <gmlpp/core/pdmat.h>

using namespace gml;

// Explicit instantiation for syntax checking

template class gml::PIsoMat<double>;
template class gml::PDiagMat<double>;
template class gml::PFullMat<double>;

/************************************************
 *
 *  Auxiliary functions
 *
 ************************************************/

struct quad_fun
{
	typedef double result_type;

	const mat& M;

	quad_fun(const mat& M_)
	: M(M_) { }

	double operator() (const colvec& x, const colvec& y) const
	{
		return arma::as_scalar(x.t() * M * y);
	}
};


template<class Derived>
bool verify_logdet(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	double v0 = math::log(arma::det(Sm));
	double v = S.log_det();

	return math::abs(v - v0) < tol;
}

template<class Derived>
bool verify_trace(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	double v0 = arma::trace(Sm);
	double v = S.trace();

	return math::abs(v - v0) < tol;
}

template<class Derived>
bool verify_mtimes(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	uword d = S.dim();
	uword n = 5;

	colvec x = arma::randn<colvec>(d);
	mat X = arma::randn<mat>(d, n);

	colvec y0 = Sm * x;
	colvec y = S * x;

	mat Y0 = Sm * X;
	mat Y = S * X;

	return is_approx(y, y0, tol) && is_approx(Y, Y0, tol);
}


template<class Derived>
bool verify_solve(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	uword d = S.dim();
	uword n = 5;

	colvec y = arma::randn<colvec>(d);
	mat Y = arma::randn<mat>(d, n);

	colvec x0 = arma::solve(Sm, y);
	colvec x1 = S.solve(y);
	colvec x2(d); S.solve(y, x2);

	mat X0 = arma::solve(Sm, Y);
	mat X1 = S.solve(Y);
	mat X2(d, n); S.solve(Y, X2);

	return is_approx(x1, x0, tol) && is_approx(x2, x0, tol) && is_approx(X1, X0, tol) && is_approx(X2, X0, tol);
}


template<class Derived>
bool verify_quad(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	uword d = S.dim();
	quad_fun f(Sm);

	colvec x = arma::randn<colvec>(d);
	colvec y = arma::randn<colvec>(d);

	double u0 = f(x, x);
	double u = S.quad(x);

	double v0 = f(x, y);
	double v = S.quad(x, y);

	return math::abs(u - u0) < tol && math::abs(v - v0) < tol;
}

template<class Derived>
bool verify_cwquad(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	uword d = S.dim();
	uword n = 5;
	quad_fun f(Sm);

	mat X = arma::randn<mat>(d, n);
	mat Y = arma::randn<mat>(d, n);

	rowvec v0 = colwise_vecmap2(f, X, Y);
	rowvec v1 = S.colwise_quad(X, Y);
	rowvec v2(n); S.colwise_quad(X, Y, v2);

	return is_approx(v1, v0, tol) && is_approx(v2, v0, tol);
}


template<class Derived>
bool verify_pwquad(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	uword d = S.dim();
	uword n = 5;
	uword n2 = 6;
	quad_fun f(Sm);

	mat X = arma::randn<mat>(d, n);
	mat Y = arma::randn<mat>(d, n2);

	mat V0 = pairwise_vecmap(f, X, Y);
	mat V1 = S.pairwise_quad(X, Y);
	mat V2(n, n2); S.pairwise_quad(X, Y, V2);

	return is_approx(V1, V0, tol) && is_approx(V2, V0, tol);
}

template<class Derived>
bool verify_inv(const PDMatBase<Derived, double>& S, const mat& Sm, double tol)
{
	mat R0 = arma::inv(Sm);

	Derived Robj = S.inv();
	mat R = Robj.get_mat();

	return is_approx(R0, R, tol);
}


/************************************************
 *
 *  Test cases
 *
 ************************************************/

TEST( PDMat, PIsoMat )
{
	// construction

	uword d = 3;
	double dv = 2.5;

	PIsoMat<double> S(d, dv);

	// basic verification

	ASSERT_EQ(d, S.dim());
	ASSERT_EQ(dv, S.diagval());

	mat Sm0 = 2.5 * arma::eye(d, d);
	mat Sm = S.get_mat();
	mat Sm2; S.get_mat(Sm2);

	ASSERT_EQ(d, Sm.n_rows);
	ASSERT_EQ(d, Sm.n_cols);
	ASSERT_EQ(d, Sm2.n_rows);
	ASSERT_EQ(d, Sm2.n_cols);
	ASSERT_TRUE( is_equal(Sm, Sm0) );
	ASSERT_TRUE( is_equal(Sm, Sm2) );

	// functional verification

	EXPECT_TRUE( verify_logdet(S, Sm, 1e-11) );
	EXPECT_TRUE( verify_trace(S, Sm, 1e-13) );
	EXPECT_TRUE( verify_mtimes(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_solve(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_quad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_cwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_pwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_inv(S, Sm, 1e-12) );
}


TEST( PDMat, PDiagMat )
{
	// construction

	const uword d = 3;
	colvec dv(d);
	dv[0] = 2.5;
	dv[1] = 1.8;
	dv[2] = 3.2;

	PDiagMat<double> S(dv);

	// basic verification

	ASSERT_EQ(d, S.dim());
	ASSERT_TRUE( is_equal(dv, S.diagvals()) );

	mat Sm0 = arma::diagmat(dv);
	mat Sm = S.get_mat();
	mat Sm2; S.get_mat(Sm2);

	ASSERT_EQ(d, Sm.n_rows);
	ASSERT_EQ(d, Sm.n_cols);
	ASSERT_EQ(d, Sm2.n_rows);
	ASSERT_EQ(d, Sm2.n_cols);
	ASSERT_TRUE( is_equal(Sm, Sm0) );
	ASSERT_TRUE( is_equal(Sm, Sm2) );

	// functional verification

	EXPECT_TRUE( verify_logdet(S, Sm, 1e-11) );
	EXPECT_TRUE( verify_trace(S, Sm, 1e-13) );
	EXPECT_TRUE( verify_mtimes(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_solve(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_quad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_cwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_pwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_inv(S, Sm, 1e-12) );
}


TEST( PDMat, PFullMat )
{
	// construction

	const uword d = 3;
	colvec ev(d);
	ev[0] = 2.5;
	ev[1] = 1.8;
	ev[2] = 3.2;

	mat H = arma::randn<mat>(d, d);
	mat Q, R;
	arma::qr(Q, R, H);

	mat C = Q.t() * arma::diagmat(ev) * Q;
	mat C0 = 0.5 * (C + C.t());
	symmetrify(C);

	ASSERT_TRUE( is_approx(C, C0, 1e-13) );

	PFullMat<double> S(C);

	// basic verification

	ASSERT_EQ(d, S.dim());
	ASSERT_TRUE( is_equal(C, S.mat()) );

	mat Sm = S.get_mat();
	mat Sm2; S.get_mat(Sm2);

	ASSERT_EQ(d, Sm.n_rows);
	ASSERT_EQ(d, Sm.n_cols);
	ASSERT_EQ(d, Sm2.n_rows);
	ASSERT_EQ(d, Sm2.n_cols);
	ASSERT_TRUE( is_equal(C, Sm) );
	ASSERT_TRUE( is_equal(C, Sm2) );

	// functional verification

	EXPECT_TRUE( verify_logdet(S, Sm, 1e-11) );
	EXPECT_TRUE( verify_trace(S, Sm, 1e-13) );
	EXPECT_TRUE( verify_mtimes(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_solve(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_quad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_cwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_pwquad(S, Sm, 1e-12) );
	EXPECT_TRUE( verify_inv(S, Sm, 1e-12) );
}





